# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.

import functools
import inspect
import logging
from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union, cast

# 配置日志
logging.basicConfig(level=logging.INFO)
logger: logging.Logger = logging.getLogger(__name__)


def _execute_operation(operation_str: Optional[str], param_dict: Dict[str, Any]) -> None:
    """执行操作代码"""
    if operation_str:
        try:
            exec(operation_str, param_dict)
        except Exception as e:
            logging.error(f"Error executing operation: {e}")


def _get_caller_info(stack: List[inspect.FrameInfo]) -> Tuple[str, str, str]:
    """获取调用者的模块名、类名、函数名"""
    caller_module: str = ""
    caller_class: str = ""
    caller_function: str = ""
    try:
        if len(stack) > 1:  # 确保有调用者栈帧
            caller_frame: inspect.FrameType = stack[1].frame
            # 模块名
            caller_module = caller_frame.f_globals.get('__name__', '')
            # 类名（处理self和cls两种情况）
            caller_self: Optional[Union[object, Type]] = caller_frame.f_locals.get('self') or caller_frame.f_locals.get(
                'cls')

            if caller_self:
                if inspect.isclass(caller_self):
                    caller_class = caller_self.__name__  # 类方法（cls）
                else:
                    caller_class = caller_self.__class__.__name__  # 实例方法（self）

            # 函数名
            caller_function = stack[1].function
    except Exception as e:
        logging.error(f"Error getting caller info: {e}")

    return caller_module, caller_class, caller_function


def _should_wrap(
        scope_name: Optional[str],
        scope_class_name: Optional[str],
        scope_function: Optional[str],
        caller_module: str,
        caller_class: str,
        caller_function: str
) -> bool:
    """判断是否需要执行包装逻辑（范围匹配）"""
    # 未指定范围时，默认执行包装
    if not all([scope_name, scope_class_name, scope_function]):
        return True

    # 检查调用者是否匹配目标范围
    return (caller_module == scope_name and
            caller_class == scope_class_name and
            caller_function == scope_function)


# 定义方法包装器类型
SyncWrapper = Callable[..., Any]
AsyncWrapper = Callable[..., Any]


def marker_prof_wrapper(
        original_method: Callable,
        params: Dict[str, Any]
) -> Union[SyncWrapper, AsyncWrapper]:
    """创建标记分析包装器"""
    function_name: Optional[str] = params.get("function_name")
    if function_name is None:
        return original_method

    entry_operation: Optional[str] = params.get("entry_operation")
    exit_operation: Optional[str] = params.get("exit_operation")
    scope_name: Optional[str] = params.get("scope_name")
    scope_class_name: Optional[str] = params.get("scope_class_name")
    scope_function: Optional[str] = params.get("scope_function")

    # 判断方法类型
    is_async: bool = inspect.iscoroutinefunction(original_method)

    if is_async:
        logger.info(f"<<<INFO: {original_method.__qualname__} is async function, use async wrapper")
        return async_func(original_method, entry_operation, exit_operation,
                          scope_name, scope_class_name, scope_function)
    else:
        logger.info(f"<<<INFO: {original_method.__qualname__} is sync function, use sync wrapper")
        return sync_func(original_method, entry_operation, exit_operation,
                         scope_name, scope_class_name, scope_function)


def sync_func(
        original_method: Callable,
        entry_operation: Optional[str],
        exit_operation: Optional[str],
        scope_name: Optional[str],
        scope_class_name: Optional[str],
        scope_function: Optional[str]
) -> SyncWrapper:
    """创建同步方法包装器"""

    @functools.wraps(original_method)
    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        # 获取当前调用栈
        _execute_operation(None, {})
        stack: List[inspect.FrameInfo] = inspect.stack()

        # 获取调用者信息
        caller_module: str
        caller_class: str
        caller_function: str
        caller_module, caller_class, caller_function = _get_caller_info(stack)

        # 判断是否需要包装
        should_wrap: bool = _should_wrap(
            scope_name, scope_class_name, scope_function,
            caller_module, caller_class, caller_function
        )

        # 如果不需要包装，直接返回原始方法
        if not should_wrap:
            logger.info(
                f"<<<INFO: Skipping wrapper for {original_method.__qualname__} "
                f"due to scope {caller_module}.{caller_class}.{caller_function} mismatch")
            return original_method(self, *args, **kwargs)

        # 准备参数字典
        param_dict: Dict[str, Any] = {
            "self": self,
            "args": args,
            "kwargs": kwargs
        }
        # 执行入口操作
        _execute_operation(entry_operation, param_dict)
        # 执行原始方法
        result: Any = original_method(self, *args, **kwargs)
        # 添加结果到参数字典
        param_dict["result"] = result
        # 执行出口操作
        _execute_operation(exit_operation, param_dict)
        return result

    return cast(SyncWrapper, wrapper)


def async_func(
        original_method: Callable,
        entry_operation: Optional[str],
        exit_operation: Optional[str],
        scope_name: Optional[str],
        scope_class_name: Optional[str],
        scope_function: Optional[str]
) -> AsyncWrapper:
    """创建异步方法包装器"""

    @functools.wraps(original_method)
    async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        # 获取当前调用栈
        _execute_operation(None, {})
        stack: List[inspect.FrameInfo] = inspect.stack()

        # 获取调用者信息
        caller_module: str
        caller_class: str
        caller_function: str
        caller_module, caller_class, caller_function = _get_caller_info(stack)

        # 判断是否需要包装
        should_wrap: bool = _should_wrap(
            scope_name, scope_class_name, scope_function,
            caller_module, caller_class, caller_function
        )

        # 如果不需要包装，直接返回原始方法
        if not should_wrap:
            logger.info(
                f"<<<INFO: Skipping wrapper for {original_method.__qualname__} "
                f"due to scope {caller_module}.{caller_class}.{caller_function} mismatch")
            return await original_method(self, *args, **kwargs)

        # 准备参数字典
        param_dict: Dict[str, Any] = {
            "self": self,
            "args": args,
            "kwargs": kwargs
        }
        # 执行入口操作
        _execute_operation(entry_operation, param_dict)
        # 执行原始方法
        result: Any = await original_method(self, *args, **kwargs)
        # 添加结果到参数字典
        param_dict["result"] = result
        # 执行出口操作
        _execute_operation(exit_operation, param_dict)
        return result

    return cast(AsyncWrapper, async_wrapper)
